from collections import deque
from typing import TYPE_CHECKING, Deque, Iterable, Iterator, List, Tuple
from chernc.treediff.gumtree.tree.fake_tree import FakeTree
from chernc.treediff.gumtree.tree.tree import TreeLike


class TreeUtils:
    @staticmethod
    def pre_order(tree: TreeLike) -> List[TreeLike]:
        trees: List[TreeLike] = []
        TreeUtils._pre_order(tree, trees)
        return trees

    @staticmethod
    def _pre_order(tree: TreeLike, trees: List[TreeLike]) -> None:
        trees.append(tree)
        if not tree.is_leaf():
            for c in tree.get_children():
                TreeUtils._pre_order(c, trees)

    @staticmethod
    def breadth_first(tree: TreeLike) -> List[TreeLike]:
        trees: List[TreeLike] = []
        currents: List[TreeLike] = [tree]
        while currents:
            c = currents.pop(0)
            trees.append(c)
            currents.extend(c.get_children())
        return trees

    @staticmethod
    def breadth_first_iterator(tree: TreeLike) -> Iterator[TreeLike]:
        class BreadthFirstIterator:
            def __init__(self):
                self.fifo: Deque[Iterator[TreeLike]] = deque()
                self._add_last(FakeTree(tree))

            def _add_last(self, item: TreeLike) -> None:
                children = item.get_children()
                if children:
                    self.fifo.append(iter(children))

            def __iter__(self) -> "BreadthFirstIterator":
                return self

            def __next__(self) -> TreeLike:
                while self.fifo:
                    it = self.fifo[0]
                    if it:
                        item = next(it)
                        if not it:
                            self.fifo.popleft()
                        self._add_last(item)
                        return item
                raise StopIteration

        return BreadthFirstIterator()

    @staticmethod
    def post_order(tree: TreeLike) -> List[TreeLike]:
        trees: List[TreeLike] = []
        TreeUtils._post_order(tree, trees)
        return trees

    @staticmethod
    def _post_order(tree: TreeLike, trees: List[TreeLike]) -> None:
        if not tree.is_leaf():
            for c in tree.get_children():
                TreeUtils._post_order(c, trees)
        trees.append(tree)

    @staticmethod
    def post_order_iterator(tree: TreeLike) -> Iterator[TreeLike]:
        class PostOrderIterator:
            def __init__(self):
                self.stack: Deque[Tuple[TreeLike, Iterator[TreeLike]]] = deque()
                self._push(tree)

            def _push(self, item: TreeLike) -> Iterator[TreeLike]:
                it = iter(item.get_children())
                self.stack.append((item, it))
                return it

            def __iter__(self) -> "PostOrderIterator":
                return self

            def __next__(self) -> TreeLike:
                if not self.stack:
                    raise StopIteration
                return self._select_next_child(self.stack[0][1])

            def _select_next_child(self, it: Iterator[TreeLike]) -> TreeLike:
                if not it:
                    return self.stack.pop()[0]
                item = next(it)
                if item.is_leaf():
                    return item
                return self._select_next_child(self._push(item))

        return PostOrderIterator()

    @staticmethod
    def pre_order_iterator(tree: TreeLike) -> Iterator[TreeLike]:
        class PreOrderIterator:
            def __init__(self):
                self.stack: Deque[Iterator[TreeLike]] = deque()
                self._push(FakeTree(tree))

            def _push(self, tree: TreeLike) -> None:
                if not tree.is_leaf():
                    self.stack.append(iter(tree.get_children()))

            def __iter__(self) -> "PreOrderIterator":
                return self

            def __next__(self) -> TreeLike:
                while self.stack:
                    it = self.stack[0]
                    if it is None:
                        raise StopIteration
                    try:
                        t = next(it)
                        if not t.is_leaf():
                            self._push(t)
                        return t
                    except StopIteration:
                        self.stack.popleft()
                raise StopIteration

        return PreOrderIterator()

    @staticmethod
    def leaf_iterator(it: Iterator[TreeLike]) -> Iterator[TreeLike]:
        class LeafIterator:
            def __init__(self, it: Iterator[TreeLike]) -> None:
                self.it = it
                self.current = next(it, None)

            def __iter__(self) -> "LeafIterator":
                return self

            def __next__(self) -> TreeLike:
                val = self.current
                while self.it:
                    self.current = next(self.it, None)
                    if self.current and self.current.is_leaf():
                        break
                if not self.it:
                    self.current = None
                return val

        return LeafIterator(it)